import numpy as np
import torch
from torchvision import datasets, transforms

from torchvision import datasets, transforms

# 定义转换操作，包括调整图像大小和转换为张量
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # 将所有图像调整为256x256大小
    transforms.ToTensor(),
])

# 加载数据集
dataset = datasets.ImageFolder(root='D:\\ai\\code\\data\\FGVC_Aircraft\\train', transform=transform)
loader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=False, num_workers=2)

# 计算均值和标准差
def compute_mean_std(loader):
    mean = 0.
    std = 0.
    total_images_count = 0

    for images, _ in loader:
        batch_samples = images.size(0)
        images = images.view(batch_samples, images.size(1), -1)
        mean += images.mean(2).sum(0)
        std += images.std(2).sum(0)
        total_images_count += batch_samples

    mean /= total_images_count
    std /= total_images_count

    return mean, std

if __name__ == '__main__':
    # 计算均值和标准差
    mean, std = compute_mean_std(loader)
    print(f'Mean: {mean}')
    print(f'Std: {std}')